package cn.doitedu.ml.demo

import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.stat.Correlation
import org.apache.spark.sql.SparkSession

import scala.collection.mutable

object CorrelationDemo {
  def main(args: Array[String]): Unit = {


    val spark = SparkSession.builder().appName("相关系数计算示例").master("local").getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._


    // val sample = spark.read.option("header", true).option("inferSchema", true).csv("user_portrait/data/correlation/sample/1.csv")
    // sample.show(100,false)
    // val arr2vec = udf((arr:mutable.WrappedArray[Double])=>{
    //   Vectors.dense(arr.toArray)
    // })
    // val vec = sample.select(arr2vec(array('面积, '房价)) as "features")
    // val res = Correlation.corr(vec, "features")
    // res.show(100,false)



    val sample = spark.read.option("header", true).option("inferSchema", true).csv("user_portrait/data/correlation/sample/2.csv")
    sample.show(100,false)
    val arr2vec = udf((arr:mutable.WrappedArray[Double])=>{
      Vectors.dense(arr.toArray)
    })
    val vec = sample.select(arr2vec(array('楼层,'面积, '房价)) as "features")

    val res = Correlation.corr(vec, "features")
    res.show(100,false)

    spark.close()
  }

}
